from sys import path
path.append('./rag')
path.append('./tools')

import os
from tools.MCPClient import MCPClient
from Agent import Agent
from rag.embedding_retriever import EmbeddingRetriever
from utils.utils import *


URL = 'https://news.ycombinator.com/'  # 未使用
out_path = os.path.join(os.getcwd(), 'output')
TASK = f"""
告诉我Antonette的信息,先从我给你的context中找到相关信息,总结后创作一个关于她的故事
把故事和她的基本信息保存到{out_path}/antonette.md,输出一个漂亮md文件
"""

fetch_mcp = MCPClient("mcp-server-fetch", "uvx", ['mcp-server-fetch'])
file_mcp = MCPClient("mcp-server-file", "npx", ['-y', '@modelcontextprotocol/server-filesystem', out_path])

async def retrieve_context():
    # RAG
    embedding_retriever = EmbeddingRetriever("BAAI/bge-m3")
    knowledge_dir = os.path.join(os.getcwd(), 'knowledge')
    files = os.listdir(knowledge_dir)
    for file in files:
        file_path = os.path.join(knowledge_dir, file)
        if os.path.isfile(file_path):
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            await embedding_retriever.embed_document(content)
    context = '\n'.join(await embedding_retriever.retrieve(TASK, 3))
    print('CONTEXT')
    print(context)
    return context

async def main():
    # RAG
    context = await retrieve_context()

    # Agent
    agent = Agent(DEFAULT_MODEL_NAME, [fetch_mcp, file_mcp], '', context)
    await agent.init()
    await agent.invoke(TASK)
    await agent.close()


if __name__ == "__main__":
    import asyncio
    asyncio.run(main())